# -*- coding: utf-8 -*-
"""
Created on Wed Apr  8 09:13:51 2020

@author: jlani
"""

import numpy as np
import pandas as pd


def mats_to_df(C,P,idnum):
    '''Pass an N x L consumption matrix C and a N x L price matrix P and an 
    idnum. 
    
    Create a dataframe with vars: cid, c1, c2, ... , cN, p1,...,pN. 
    The indices will be idnum-obnum.'''
    
    N, L = np.shape(C)
    
    consumerid_list = [ ['cid', idnum] ]
    #Make a list of lists where each sublist is ['c(n+1)', C[n]]
    c_list = [[ 'c' + str(l+1), C[:,l]] for l in range(L) ]
    #Same list of lists for prices.
    p_list = [[ 'p' + str(l+1), P[:,l]] for l in range(L) ]
    df_dict = dict( consumerid_list + c_list + p_list ) #dict to make dataframe.
    indices = [ str(idnum) + '-' + str(n) for n in range(N) ]
    return pd.DataFrame( df_dict, indices )
    

def df_to_mats(df,idnum):
    '''Pass a dataframe which has variables:
    cid, c1, c2, ..., cL, p1, p2, ..., pL.
        
    Will create a consumption matrix with c1,...,cL and a price matrix with
    p1,...,pL using data for cid == idnum. 
    Return as a list with two elements [C,P].
    '''
    
    #Get the consumption and price variable names.
    c_name_list = []; p_name_list = []
    for varname in df.columns:
        if varname[0] == 'c':
            try: 
                cnum = int(varname[1:]) #Will raise ValueError if can't make into an integer.
                c_name_list.append('c' + str(cnum) )
                p_name_list.append('p' + str(cnum) )        
            except ValueError:
                pass
    obs_to_use = df[ 'cid' ] == idnum    
    C = np.array( df[c_name_list][ obs_to_use ] )
    P = np.array( df[p_name_list][ obs_to_use ] )
    return [C,P]
    
    
    



